#pragma once

#include "common_unit.h"
#include <queue>
#include <thread>
#include <mutex>
#include <memory>
#include "doca_preprocess_unit.h"
#include "debug_unit.h"
#include "timer.h"
#include "queue_unit.h"
#include "memory_unit.h"
#include "doca_brpc_unit.h"
#define MAX_QUEUE_SIZE 8

class ThreadSafeQueue {
public:
    ThreadSafeQueue(size_t pool_size, const std::vector<std::shared_ptr<MemoryNode> >& memnode_list)
        : pool_size_(pool_size), head_(0), tail_(0), count_(0) {
        
        memory_pool_.resize(pool_size_);
        memnode_size.clear();

        block_size_ = 0;
        for (auto mem_node : memnode_list) {
            memnode_size.push_back(mem_node->total_size);
            block_size_ += mem_node->total_size;
        }
        for (size_t i = 0; i < pool_size_; ++i) {
            memory_pool_[i] = new char[block_size_];
        }
    }

    ~ThreadSafeQueue() {
        for (auto ptr : memory_pool_) {
            delete[] reinterpret_cast<char*>(ptr);
        }
    }

    bool push(const std::vector<void*>& data) {
        std::lock_guard<std::mutex> lock(mutex_);
        if (count_ >= pool_size_) {
            return false; // Queue is full
        }
        if ( (int)data.size() != (int)memnode_size.size()) {
            printf ("memcpy size not match!\n");
            return false;
        }

        size_t addr_start = 0;
        for (int i = 0; i < (int)memnode_size.size(); ++i) {
            std::memcpy(memory_pool_[tail_] + addr_start, data[i], memnode_size[i]);
            addr_start += memnode_size[i];
        }
        tail_ = (tail_ + 1) % pool_size_;
        count_++;
        return true;
    }

    bool pop(std::vector<void*>& data) {
        std::lock_guard<std::mutex> lock(mutex_);
        if (count_ == 0) {
            return false; // Queue is empty
        }
        if (data.size() != memnode_size.size()) {
            std::cerr << "memcpy size not match!" << std::endl;
            return false;
        }

        size_t addr_start = 0;
        for (size_t i = 0; i < memnode_size.size(); ++i) {
            std::memcpy(data[i], memory_pool_[head_] + addr_start, memnode_size[i]);
            addr_start += memnode_size[i];
        }
        head_ = (head_ + 1) % pool_size_;
        count_--;
        return true;
    }

private:
    std::vector<void*> memory_pool_;
    std::vector<size_t> memnode_size;
    size_t pool_size_;
    size_t block_size_;
    size_t head_;
    size_t tail_;
    size_t count_;
    std::mutex mutex_;
};

class StageUnit{
public:
    int stage_id;

    std::vector<int> input_mem_id_list;
    std::vector<int> output_mem_id_list;

    std::vector<std::shared_ptr<MemoryNode> > input_mem_node_list;
    std::vector<std::shared_ptr<MemoryNode> > output_mem_node_list;

    std::vector<std::shared_ptr<OperatorNode> > subdag;
    std::vector<std::shared_ptr<OperatorNode> > sequencial_list;
    std::map<int , void *> mem_addr_map;
    std::map<int , bool> visit_op;
    std::map<int , std::shared_ptr<MemoryNode> > id_to_mem_node;

    std::shared_ptr<QueueUnit> input_pipe_queue;
    std::shared_ptr<QueueUnit> output_pipe_queue;
    std::shared_ptr<QueueUnit> output_mem_queue;
    std::shared_ptr<QueueUnit> recycle_mem_queue;

    void init_subdag(std::vector<std::shared_ptr<OperatorNode> > & subdag) {this->subdag = subdag;}
    void mem_init();
    void sequence_assign();
    void execute_dag();
    void execute_pipeline();

    virtual void mem_assign() = 0;
    virtual void init_processunit() = 0;
    
    bool exit_flag = false;
    bool if_last_stage = false;
    bool if_first_stage = false;
    stage_type_t stage_type = stage_type_t::NORMAL_STAGE;
    size_t total_size = 0;
public:
    PreProcessUnit* processunit;
};

class CpuStageUnit : public StageUnit {
public:
    virtual void mem_assign() override;
    virtual void init_processunit() override;
};

class DpaStageUnit : public StageUnit {
public:
    virtual void mem_assign() override;
    virtual void init_processunit() override;
public:
    void * DpaBaseAddr;
};

#define MAX_STAGE_NUM 3

class PipelineUnit{
public:
    std::vector<StageUnit *> stage_list[MAX_STAGE_NUM];
public:
    void Init_stage_list(
        std::vector<int> stage_thread_num, 
        std::vector<std::vector<std::shared_ptr<OperatorNode> > > & subdag_list,
        std::vector<device_type_t> device_type_list,
        std::vector<stage_type_t> stage_type_list,
        int if_cross_device
    );
    void start_pipeline();
    void end_pipeline();
    std::vector<std::thread> worker_threads; 
};

void StageUnit::mem_init() {
    for (auto op_node : subdag) {
        visit_op[op_node->id] = true;
    }

    this->total_size = 0;
    
    for (auto op_node : subdag) {
        bool if_rd_zero = true;
        bool if_cd_zero = true;

        for (auto next_mem_node : op_node->next_mem_nodes) {
            this->total_size += next_mem_node->total_size;
        }

        for (auto neighbor : op_node->prev_ops) {
            if (visit_op[neighbor->id]) {
                if_rd_zero = false;
                break;
            }
        }
        if (if_rd_zero) {
            for (auto prev_mem_node : op_node->prev_mem_nodes) {
                input_mem_id_list.push_back(prev_mem_node->id);
                this->id_to_mem_node[prev_mem_node->id] = prev_mem_node;
            }
        }

        for (auto neighbor : op_node->next_ops) {
            if (visit_op[neighbor->id]) {
                if_cd_zero = false;
                break;
            }
        }

        if (if_cd_zero) {
            // TODO: 解决op_node直接就是全图出度为0的节点的情况，可以考虑增加几个Stage这样，做重复操作。

            if (this->stage_type != stage_type_t::SEND_STAGE) { // last_stage only for normal_stage type
                if (this->if_last_stage == false && (int)op_node->next_ops.size() == 0) {
                    continue;
                }
            }

            else { // send_stage means not last_stage
                if ((int)op_node->next_ops.size() == 0) continue;
            }

            for (auto next_mem_node : op_node->next_mem_nodes) {
                output_mem_id_list.push_back(next_mem_node->id);
                output_mem_node_list.push_back(next_mem_node);
            }
        }
    }

    std::sort(input_mem_id_list.begin(), input_mem_id_list.end() );
    auto last = std::unique(input_mem_id_list.begin(), input_mem_id_list.end());
    input_mem_id_list.erase(last, input_mem_id_list.end());

    for (auto mem_id : input_mem_id_list) {
        this->input_mem_node_list.push_back(this->id_to_mem_node[mem_id]);
        this->total_size += this->id_to_mem_node[mem_id]->total_size;
    }

    std::sort(output_mem_id_list.begin(), output_mem_id_list.end() );
    std::sort(output_mem_node_list.begin(), output_mem_node_list.end(), CommonUtils::tmp_mem_node);
}

void StageUnit::sequence_assign() {
    std::map<int, int> in_degree; // 用于存储每个节点的入度
    std::queue<std::shared_ptr<OperatorNode>> zero_in_degree_queue; // 存储入度为零的节点

    // 初始化入度，仅计算属于 subdag 的节点
    for (auto op_node : subdag) {
        in_degree[op_node->id] = 0;
    }
    for (auto op_node : subdag) {
        for (auto neighbor : op_node->next_ops) {
            if (visit_op[neighbor->id]) { // 只考虑属于 subdag 的后继节点
                in_degree[neighbor->id]++;
            }
        }
    }

    // 将入度为零的节点加入队列
    for (auto op_node : subdag) {
        if (in_degree[op_node->id] == 0) {
            zero_in_degree_queue.push(op_node);
        }
    }

    // 拓扑排序
    while (!zero_in_degree_queue.empty()) {
        auto current_node = zero_in_degree_queue.front();
        zero_in_degree_queue.pop();
        
        sequencial_list.push_back(current_node); // 加入顺序列表

        // 更新后继节点的入度
        for (auto neighbor : current_node->next_ops) {
            if (visit_op[neighbor->id]) { // 只考虑属于 subdag 的后继节点
                in_degree[neighbor->id]--;
                if (in_degree[neighbor->id] == 0) {
                    zero_in_degree_queue.push(neighbor); // 入度为零，加入队列
                }
            }
        }
    }

    // 检查是否所有节点都被排序（如果不是，说明存在环）
    if (sequencial_list.size() != subdag.size()) {
        throw std::runtime_error("Graph contains a cycle, cannot perform topological sort.");
    }
}

void CpuStageUnit::init_processunit() {
    processunit = new CpuPreProcessUnit();
    processunit->thread_init(1);
}

void CpuStageUnit::mem_assign() {
    MemoryUnit & memoryunit = MemoryUnit::Instance();
    for (auto op_node : subdag) {
        for (auto neighbor : op_node->prev_mem_nodes) {
            if (!mem_addr_map[neighbor->id]) mem_addr_map[neighbor->id] = memoryunit.Malloc(neighbor->total_size);
        }
        for (auto neighbor : op_node->next_mem_nodes) {
            if (!mem_addr_map[neighbor->id]) mem_addr_map[neighbor->id] = memoryunit.Malloc(neighbor->total_size);
        }
    }
}

void StageUnit::execute_dag() {
    processunit->task_execute(sequencial_list, this->mem_addr_map);
}

void StageUnit::execute_pipeline() {
    xmh::Timer stage_timer("stage_" + std::to_string(this->stage_id) + "_total_time");
    xmh::Timer stage_execute_timer("stage_" + std::to_string(this->stage_id) + "_per_execute_time" );

    int dag_count = 0;
    stage_timer.start();

    while (true && exit_flag == false) {
        // 从输入队列中获取input的指针（from pipe）和output的指针（from mem）
        void * input_ptr;
        void * output_ptr;

        while (exit_flag == false && !output_mem_queue->pop(output_ptr)) {}
        while (exit_flag == false && !input_pipe_queue->pop(input_ptr) ) {}

        /* assign input and output */
        size_t total_used_space = 0;
        for (auto Mem_node : this->input_mem_node_list) {
            this->mem_addr_map[Mem_node->id] = input_ptr + total_used_space;
            total_used_space += Mem_node->total_size;
        }

        total_used_space = 0;
        for (auto Mem_node : this->output_mem_node_list) {
            this->mem_addr_map[Mem_node->id] = output_ptr + total_used_space;
            total_used_space += Mem_node->total_size;
        }

        stage_execute_timer.start();
        this->execute_dag();
        stage_execute_timer.end();
        // debug
        ++ dag_count;
        if (this->if_last_stage == true && this->stage_type != stage_type_t::SEND_STAGE) {
            if (dag_count == 500) {
                for (auto mem_node : output_mem_node_list) {
                    DebugUnit::mem_debug(mem_node, 10, mem_addr_map[mem_node->id]);
                }
            }
        }

        // to output_queue;
        while (exit_flag == false && !recycle_mem_queue->push(input_ptr)  ) {}
        while (exit_flag == false && !output_pipe_queue->push(output_ptr) ) {}
    }

    LOG(INFO) << "execute_batch " << dag_count << " times";
    stage_timer.end();
}

void DpaStageUnit::init_processunit() {
    processunit = new DpaPreProcessUnit();
}

void DpaStageUnit::mem_assign() {
    MemoryUnit & memoryunit = MemoryUnit::Instance();
    this->DpaBaseAddr = memoryunit.ptr;
    size_t total_size = memoryunit.total_size;
    DpaPreProcessUnit* dpaProcessUnit = static_cast<DpaPreProcessUnit*>(processunit);
    dpaProcessUnit->register_host_memory(this->DpaBaseAddr, total_size);
    processunit->thread_init(16); //到这里才能init

    for (auto op_node : subdag) {
        for (auto neighbor : op_node->prev_mem_nodes) {
            if (!mem_addr_map[neighbor->id]) mem_addr_map[neighbor->id] = memoryunit.Malloc(neighbor->total_size);
        }
        for (auto neighbor : op_node->next_mem_nodes) {
            if (!mem_addr_map[neighbor->id]) mem_addr_map[neighbor->id] = memoryunit.Malloc(neighbor->total_size);
        }
    }

    LOG(INFO) << "start init dpa thread info";
    dpaProcessUnit->init_dpa_thread_info(this->subdag, &mem_addr_map);
    LOG(INFO) << "finish init dpa thread info";
}

void PipelineUnit::Init_stage_list(
        std::vector<int> stage_thread_num, 
        std::vector<std::vector<std::shared_ptr<OperatorNode> > > & subdag_list,
        std::vector<device_type_t> device_type_list,
        std::vector<stage_type_t> stage_type_list,
        int if_cross_device
    ) {
    /* 首先新建若干stage，把input_mem_list和output_mem_list都建好 */
    int stage_num = (int)stage_thread_num.size();
    for (int i = 0; i < (int)stage_num; ++i) {
        for (int j = 0; j < stage_thread_num[i]; ++j) {
            //TODO: 每个阶段用不同类型的StageUnit
            StageUnit * cur_unit;
            if (device_type_list[i] != device_type_t::DEVICE_DPA) {
                cur_unit = new CpuStageUnit();
            }
            else{
                cur_unit = new DpaStageUnit();
            }

            if (i == 0) cur_unit->if_first_stage = true;
            if (i == (int)stage_thread_num.size() - 1) cur_unit->if_last_stage = true;
            cur_unit->stage_type = stage_type_list[i];

            cur_unit->stage_id = i;
            cur_unit->init_processunit();
            cur_unit->init_subdag(subdag_list[i]);
            cur_unit->mem_init();
            // cur_unit->mem_assign();
            cur_unit->sequence_assign();
            stage_list[i].push_back(cur_unit);
        }
    }

    /* 然后建立stage_num + 1个MemQueueUnit，统计内存信息，接着初始化内存（初始化MemQueueUnit） */
    std::vector<std::shared_ptr<MemQueueUnit> > MemQueueList;
    size_t total_used_space = 0;
    for (int i = 0; i < stage_num; ++i) {
        total_used_space += this->stage_list[i][0]->total_size * (int)stage_list[i].size();

        if (i == 0) {
            MemQueueList.push_back(std::make_shared<MemQueueUnit>(MAX_QUEUE_SIZE, stage_list[i][0]->input_mem_node_list) );
        }

        MemQueueList.push_back(std::make_shared<MemQueueUnit>(MAX_QUEUE_SIZE, stage_list[i][0]->output_mem_node_list) );
        /* 默认会设置成NORMAL_STAGE */
        if (this->stage_list[i][0]->stage_type == stage_type_t::SEND_STAGE) {
            MemQueueList[i+1]->stage_type = stage_type_t::SEND_STAGE;
        }
        else if (this->stage_list[i][0]->stage_type == stage_type_t::RECV_STAGE) {
            MemQueueList[i]->stage_type = stage_type_t::RECV_STAGE;
        }
    }

    for (auto MemQueue : MemQueueList) {
        total_used_space += MemQueue->get_total_mem_size();
    }

    /* 统一管理内存，开一块大的 */
    MemoryUnit & memoryunit = MemoryUnit::Instance();
    memoryunit.Init(total_used_space);

    // assign mem

    for (int i = 0; i < stage_num; ++i) {
        for (auto & stage : stage_list[i]) {
            stage->mem_assign();
        }
    }

    for (auto MemQueue : MemQueueList) {
        if (MemQueue->stage_type == stage_type_t::NORMAL_STAGE) MemQueue->mem_assign();
    }

    /* start dma info initialize */
    void * Dma_start_ptr = nullptr;
    std::vector<size_t> Dma_offsets;

    size_t acc_offset = 0;
    Dma_offsets.push_back(acc_offset);

    for (auto MemQueue : MemQueueList) {
        if (MemQueue->stage_type != stage_type_t::NORMAL_STAGE) {
            MemQueue->mem_assign();

            size_t cur_mem_size = MemQueue->get_mem_size();
            while (cur_mem_size > (1 << 21) ) {

                MemQueue->Mem_task_id.push_back( (int)Dma_offsets.size() - 1);

                acc_offset += (1 << 21);
                Dma_offsets.push_back(acc_offset);

                cur_mem_size -= (1 << 21);
            }

            if (cur_mem_size > 0) {
                MemQueue->Mem_task_id.push_back( (int)Dma_offsets.size() - 1);
                acc_offset += cur_mem_size;
                Dma_offsets.push_back(acc_offset);
            }

            MemQueue->SendFlag_id = (int)Dma_offsets.size() - 1;
            acc_offset += sizeof(int64_t);
            Dma_offsets.push_back(acc_offset);
            
            MemQueue->RecvFlag_id = (int)Dma_offsets.size() - 1;
            acc_offset += sizeof(int64_t);
            Dma_offsets.push_back(acc_offset);
            
        }
    }

    /* 可能有多个DMA的stage（可扩展性），这里统一分配内存，第一个start的ptr即为开头 */
    for (auto MemQueue : MemQueueList) {
        if (MemQueue->stage_type != stage_type_t::NORMAL_STAGE) {
            Dma_start_ptr = MemQueue->MemoryPool[0];
            break;
        }
    }

    /* DMA register */
    if (if_cross_device) {
        // check
        if (Dma_start_ptr == nullptr) LOG(INFO) << "Dma addr pointer is null !!";
        BrpcUnit & brpcunit = BrpcUnit::GetInstance();
        brpcunit.InitDmaUnit(Dma_offsets, Dma_start_ptr);
#if defined(__x86_64__) // for cpu, start brpc, send export info
        brpcunit.InitBrpc();
        brpcunit.SendDmaInfo();
#endif
    }

    // init MemQueue[0];
    if (MemQueueList[0]->stage_type != stage_type_t::RECV_STAGE) {
        MemQueueList[0]->Mem_Init();
    }

    /* 建立stage_num + 1个PipeQueueUnit， */

    std::vector<std::shared_ptr<PipeQueueUnit> > PipeQueueList;

    for (int i = 0; i < stage_num + 1; ++i) {
        bool if_first_stage = (i == 0);
        bool if_last_stage  = (i == stage_num);

        PipeQueueList.push_back(std::make_shared<PipeQueueUnit>(MAX_QUEUE_SIZE, MemQueueList[i], if_first_stage, if_last_stage) );
    }

    /* 初始化StageUnit的所有要用到的queue */
    for (int i = 0; i < stage_num; ++i) {
        for (auto & stage : stage_list[i]) {
            stage->input_pipe_queue  = PipeQueueList[i];
            stage->output_pipe_queue = PipeQueueList[i+1];
            stage->output_mem_queue  = MemQueueList[i+1];
            stage->recycle_mem_queue = MemQueueList[i];
        }
    }
}

void PipelineUnit::start_pipeline(){
    xmh::Reporter::StartReportThread();
    for (int i = 0; i < MAX_STAGE_NUM; ++i) {
        if ( (int)stage_list[i].size() == 0) break;
        for (StageUnit* stage_unit : stage_list[i]) {
            // Create a thread for each StageUnit's execute method
            worker_threads.emplace_back([stage_unit]() {
                stage_unit->execute_pipeline();
            });
        }
    }
}

void PipelineUnit::end_pipeline() {
    for (int i = 0; i < MAX_STAGE_NUM; ++i) {
        if ( (int)stage_list[i].size() == 0) break;
        for (StageUnit* stage_unit : stage_list[i]) {
            stage_unit->exit_flag = true;
        }
    }

    for (std::thread& worker_thread : worker_threads) {
        if (worker_thread.joinable()) {
            worker_thread.join();
        }
    }

    xmh::Reporter::StopReportThread();
    xmh::Reporter::Report();
}